
#include "div_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

constexpr uint32_t BLOCK_SIZE = 32;
constexpr int32_t inputVarNum = 2;
constexpr int32_t maxDimNum = 64;

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    DivTilingData tiling;
    const gert::StorageShape* x1_shape = context->GetInputShape(0);
    int32_t totallength = 1;
    for (int i = 0; i < x1_shape->GetStorageShape().GetDimNum(); i++)
        totallength *= x1_shape->GetStorageShape().GetDim(i);
    tiling.set_totallength(totallength);

    uint64_t ubSize;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize(); 
    uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(0)->GetDataType()); 
    uint32_t inputLength = inputBytes * inputNum;
    uint32_t ubDataNumber = (inputBytes == 2) ? 12 : 12;
    uint32_t tileBlockNum = (ubSize / BLOCK_SIZE) / ubDataNumber; //每个ub段可用的空间块数
    uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes; //每次处理的数据量
    tiling.set_tileDataNum(tileDataNum);
    context->SetTilingKey(1);

    uint32_t x1_length = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
    uint32_t x2_length = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
    int64_t shape[maxDimNum * inputVarNum], shapefull[maxDimNum];
    int64_t numshapes = 1;
    if(x1_length != x2_length) {
        context->SetTilingKey(2);
        int64_t numshapes1 = context->GetInputShape(0)->GetStorageShape().GetDimNum();
        int64_t numshapes2 = context->GetInputShape(1)->GetStorageShape().GetDimNum();
        numshapes = (numshapes1 > numshapes2) ? numshapes1 : numshapes2;
        for (int k = 0; k < inputVarNum; ++k) {
            int64_t *ss = &shape[k * maxDimNum];
            const gert::StorageShape* inputshape = context->GetInputShape(k);
            for (int i = 0; i < inputshape->GetStorageShape().GetDimNum(); i++) {
                ss[i] = inputshape->GetStorageShape().GetDim(i);
            }
        }
        for (int k = 0; k < numshapes; ++k) {
            int64_t *ss = &shape[0];
            int64_t *sf = &shapefull[0];
            sf[k] = (ss[k] > ss[k + maxDimNum]) ? ss[k] : ss[k + maxDimNum];   
        }
    }
    tiling.set_x1_length(x1_length);
    tiling.set_x2_length(x2_length);
    tiling.set_numshapes(numshapes);
    tiling.set_shape(shape);
    tiling.set_shapefull(shapefull);
    std::cout << "tiling_key:" << context->GetTilingKey() << std::endl;
    std::cout << "totallength:" << totallength << std::endl;
    std::cout << "x1_length:" << x1_length << std::endl;
    std::cout << "x2_length:" << x2_length << std::endl;
    std::cout << "tileDataNum:" << tileDataNum << std::endl;
    std::cout << "numshapes:" << numshapes << std::endl;

    context->SetBlockDim(1);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Div : public OpDef {
public:
    explicit Div(const char* name) : OpDef(name)
    {
        this->Input("x1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(Div);
}
